import cv2
import numpy as np
import math

def read_images(p1_path, p2_path):
    """读取原图P1和待嵌入图P2，预处理为统一尺寸的灰度图"""
    p1 = cv2.imread(p1_path, cv2.IMREAD_GRAYSCALE)
    p2 = cv2.imread(p2_path, cv2.IMREAD_GRAYSCALE)
    
    if p1 is None:
        raise ValueError(f"原图P1读取失败，请检查路径：{p1_path}")
    if p2 is None:
        raise ValueError(f"待嵌入图P2读取失败，请检查路径：{p2_path}")
    
    # 统一尺寸（以P1为基准）
    p2_resized = cv2.resize(p2, (p1.shape[1], p1.shape[0]))
    # P2二值化，适配嵌入需求
    _, p2_binary = cv2.threshold(p2_resized, 127, 1, cv2.THRESH_BINARY)
    return p1, p2_binary

def embed_image(p1, p2_binary):
    """嵌入P2到P1，生成隐写图P"""
    stego_image = p1.copy()
    height, width = stego_image.shape
    for i in range(height):
        for j in range(width):
            stego_image[i, j] = (stego_image[i, j] >> 1) << 1 | p2_binary[i, j]
    return stego_image

def calculate_psnr(p1, stego_image):
    """计算PSNR值，验证视觉相似度"""
    mse = np.mean((p1.astype(np.float32) - stego_image.astype(np.float32)) ** 2)
    if mse == 0:
        return float('inf')
    psnr = 10 * math.log10((255 ** 2) / mse)
    return psnr

def save_image(image, save_path):
    """保存隐写图"""
    cv2.imwrite(save_path, image)
    print(f"隐写图已保存至：{save_path}")

def main():
    # 已替换为实际图片格式路径
    P1_PATH = "D:\\CodeFiles\\Gitee\\Python\\BigHomework\\Algorithm\\bigwork.jpg"    # P1实际路径（JPG格式）
    P2_PATH = "D:\\CodeFiles\\Gitee\\Python\\BigHomework\\Algorithm\\bigwork01.png"  # P2实际路径（PNG格式）
    SAVE_PATH = "D:\\CodeFiles\\Gitee\\Python\\BigHomework\\Algorithm\\stego_p.jpg"  # 隐写图保存路径
    
    try:
        p1, p2_binary = read_images(P1_PATH, P2_PATH)
        stego_image = embed_image(p1, p2_binary)
        save_image(stego_image, SAVE_PATH)
        psnr_value = calculate_psnr(p1, stego_image)
        print(f"原图与隐写图的PSNR值：{psnr_value:.2f} dB")
        print("隐写完成！人眼视觉上隐写图与原图基本一致")
    except Exception as e:
        print(f"程序执行出错：{str(e)}")

if __name__ == "__main__":
    main()